import pickle as pkl
import joblib
import os
import numpy as np
os.environ["CUDA_VISIBLE_DEVICES"] = "4"
class Query:
    def __init__(self, i):
        pos_feat_file =None
        pos_pcode_seq_file =None
        pos_artid_seq_file=None
        if i < 31:
            pos_feat_file = open(f'./Data/Processed/divide/pos_feat{i}.db','rb')
            pos_pcode_seq_file = open(f'./Data/Processed/divide/pos_pcode_seq{i}.db','rb')
            pos_artid_seq_file = open(f'./Data/Processed/divide/pos_artid_seq{i}.db','rb')
        else:
            pos_feat_file = open(f'./Data/Processed/divide/pos_feat_val.db','rb')
            pos_pcode_seq_file = open(f'./Data/Processed/divide/pos_pcode_seq_val.db','rb')
            pos_artid_seq_file = open(f'./Data/Processed/divide/pos_artid_seq_val.db','rb')

        self.pos_custom_id = joblib.load(pos_feat_file)
        self.pos_FN = joblib.load(pos_feat_file)
        self.pos_Active = joblib.load(pos_feat_file)
        self.pos_club_member_status = joblib.load(pos_feat_file)
        self.pos_fashion_news_frequency = joblib.load(pos_feat_file)
        self.pos_age = joblib.load(pos_feat_file)
        self.pos_artid = joblib.load(pos_feat_file)
        self.pos_pcode = joblib.load(pos_feat_file)
        self.pos_sales_channel_id = joblib.load(pos_feat_file)
        
        self.pos_artid_seq = joblib.load(pos_artid_seq_file)
        self.pos_pcode_seq = joblib.load(pos_pcode_seq_file)
        self.index = np.arange(self.pos_artid.size)
        np.random.shuffle(self.index)
        self.pos_custom_id = self.pos_custom_id[self.index]
        self.pos_FN = self.pos_FN[self.index]
        self.pos_Active = self.pos_Active[self.index]
        self.pos_club_member_status = self.pos_club_member_status[self.index]
        self.pos_fashion_news_frequency = self.pos_fashion_news_frequency[self.index]
        self.pos_age = self.pos_age[self.index]
        self.pos_artid = self.pos_artid[self.index]
        self.pos_pcode = self.pos_pcode[self.index]
        self.pos_sales_channel_id = self.pos_sales_channel_id[self.index]
        self.pos_artid_seq = self.pos_artid_seq[self.index,...]
        self.pos_pcode_seq = self.pos_pcode_seq[self.index,...]
        #print(type(self.pos_pcode_seq))
        print('load ok')
        self.start = 0

    def query(self, N):
        if self.start + N > len(self.pos_custom_id):
            N = len(self.pos_custom_id) - self.start
        pos_custom_id = self.pos_custom_id[self.start:self.start+N]
        pos_FN = self.pos_FN[self.start:self.start+N]
        pos_Active = self.pos_Active[self.start:self.start+N]
        pos_club_member_status = self.pos_club_member_status[self.start:self.start+N]
        pos_fashion_news_frequency = self.pos_fashion_news_frequency[self.start:self.start+N]
        pos_age = self.pos_age[self.start:self.start+N]
        pos_artid = self.pos_artid[self.start:self.start+N]
        pos_pcode = self.pos_pcode[self.start:self.start+N]
        pos_sales_channel_id = self.pos_sales_channel_id[self.start:self.start+N]
        pos_pcode_seq = self.pos_pcode_seq[self.start:self.start+N]
        pos_artid_seq = self.pos_artid_seq[self.start:self.start+N]
        seq_length = np.array([i for i in np.count_nonzero(pos_artid_seq, axis = 1)])
        pos_click = np.array([1 for i in range(N)])
        self.start += N
        return {"pos_custom_id":pos_custom_id, \
        "pos_FN":pos_FN, \
        "pos_Active":pos_Active, \
        "pos_club_member_status":pos_club_member_status,\
        "pos_fashion_news_frequency":pos_fashion_news_frequency, \
        "pos_age":pos_age, \
        "pos_artid":pos_artid, \
        "pos_pcode":pos_pcode, \
        "pos_sales_channel_id":pos_sales_channel_id, \
        "hist_pos_pcode":pos_pcode_seq,\
        "hist_pos_artid":pos_artid_seq,\
        "seq_length":seq_length,\
        "pos_click":pos_click,\
        "length":N}

